#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.4
# @author : cd
# @Email : 19688513@qq.com
# @Project : new-horizons-engine
# @File : SignalGenerator.py
# @Software: PyCharm
# @time: 2025/6/4 10:18
# @description : 信号生成器类（多线程优化版）
"""

import logging
import pandas as pd
import numpy as np
import sqlite3
from datetime import datetime
from typing import Dict, Any, List, Optional, Union
import threading
import os
import re
from concurrent.futures import ThreadPoolExecutor, as_completed

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


# 改进的成交量信号条件
def calculate_volume_signal(signals):

    # 1. 当前成交量接近5日均量（容差1%）
    volume_near_ma5 = (abs(signals['volume'] - signals['vol_ma5']) /
                       signals['vol_ma5'].clip(lower=1e-6) < 0.01)

    # 2. 5日均量低于120日均量（缩量趋势）
    ma5_below_ma120 = signals['vol_ma5'] < signals['vol_ma120']

    # 3. 当前成交量明显低于60日均量（缩量确认）
    # 使用动态阈值：基于市场波动性调整
    volume_volatility = signals['volume'].pct_change().rolling(window=20).std().fillna(0)
    dynamic_threshold = 0.8 + 0.1 * (volume_volatility / 0.02)  # 波动大时放宽条件
    dynamic_threshold = np.clip(dynamic_threshold, 0.75, 0.85)  # 限制在合理范围

    volume_below_ma60 = signals['volume'] < signals['vol_ma60'] * dynamic_threshold

    # 4. 增加成交量形态确认（连续3日成交量低于120日均量）
    low_volume_pattern = (signals['volume'] < signals['vol_ma120']).rolling(window=3).mean() >= 0.67

    # 5. 结合所有条件
    signals['signal_vol'] = (
            volume_near_ma5 &
            ma5_below_ma120 &
            volume_below_ma60 &
            low_volume_pattern
    ).fillna(False)

    return signals


class SignalGenerator:
    """信号生成器类，负责生成交易信号（多线程优化版）"""

    def __init__(self, cache_db_path: str = "data/signal_cache.db", max_workers: int = 4):
        # 初始化缓存数据库
        self.cache_db_path = cache_db_path
        self.max_workers = max_workers
        self.local = threading.local()  # 先初始化线程本地存储

        # 信号名称和编号的映射 - 根据成交量重新分类
        self.SIGNAL_MAP = {
            # 成交量信号
            'signal_vol': 1,  # 成交量双均线条件
            'signal_vol_break': 2,  # 量能趋势突破
            'signal_volume_price': 3,  # 量价齐升
            'signal_vol_macd': 4,  # 成交量MACD金叉
            'signal_vol_rsi': 5,  # 成交量RSI超卖
            'signal_vol_boll': 6,  # 成交量布林带收窄
            'signal_vol_kdj': 7,  # 成交量KDJ金叉
            'signal_low_volatility': 8,  # 低波动率信号

            # 价格信号
            'signal_price': 11,  # 价格均线多头排列
            'signal_trend': 12,  # 趋势延续验证
            'signal_volatility': 13,  # 波动率适中
            'signal_rsi': 14,  # 价格RSI超卖

            # 其他信号
            'signal_support_test': 21,  # 关键支撑测试
            'signal_breakout': 22,  # 突破信号
            'signal_reversal': 23  # 反转信号
        }

        # 信号描述映射 - 根据成交量重新组织
        self.SIGNAL_DESCRIPTIONS = {
            # 成交量信号描述
            1: "复合成交量信号：当前成交量接近5日均量(±1%)，5日均量低于120日均量，当前成交量低于60日均量的动态阈值(75%-85%)，且近3日中有至少2日成交量低于120日均量",
            2: "量能趋势突破：5日成交量均线上穿20日成交量均线，且5日均量近期呈上升趋势",
            3: "量价齐升：成交量放大10%以上，阳线且突破前日高点",
            4: "成交量MACD金叉：成交量MACD上穿信号线",
            5: "成交量RSI超卖：成交量RSI < 30",
            6: "成交量布林带收窄：带宽小于20日均宽的80%，且成交量在通道内",
            7: "成交量KDJ金叉：K线上穿D线且J值>20",
            8: "低波动率信号：连续多日成交量低于20日均量的80%，且价格波动率低于近期平均水平",

            # 价格信号描述
            11: "价格均线多头排列：5日>10日>20日均线，且5日均线上扬",
            12: "趋势延续验证：成交量均线多头排列(5>20>60)且价格高于3日均价",
            13: "波动率适中：14日ATR处于近1年30%-70%分位之间",
            14: "价格RSI超卖：价格RSI低于其14日均值减1.5倍标准差",

            # 其他信号描述
            21: "关键支撑测试：价格接近近期低点且成交量萎缩",
            22: "突破信号：价格突破20日高点且成交量放大",
            23: "反转信号：长下影线且成交量放大"
        }

        # 确保数据库目录存在
        db_dir = os.path.dirname(cache_db_path)
        if not os.path.exists(db_dir):
            os.makedirs(db_dir)  # 创建目录

        # 初始化数据库
        self._init_cache_database()

    def _get_db_connection(self):
        """获取当前线程的数据库连接"""
        # 如果当前线程没有连接，创建一个新的
        if not hasattr(self.local, 'conn') or not self.local.conn:
            self.local.conn = sqlite3.connect(self.cache_db_path)
            self.local.conn.execute("PRAGMA journal_mode = WAL")
            self.local.conn.execute("PRAGMA synchronous = NORMAL")
            self.local.conn.execute("PRAGMA cache_size = -20000")  # 20MB缓存
        return self.local.conn

    def _init_cache_database(self):
        """初始化缓存数据库结构"""
        # 创建信号缓存表
        self._create_signal_cache_table()

        # 为每个信号创建独立的缓存表
        for signal_id in self.SIGNAL_MAP.values():
            self._create_signal_table(signal_id)

    def _create_signal_cache_table(self):
        """创建信号缓存表"""
        conn = self._get_db_connection()
        conn.execute('''
            CREATE TABLE IF NOT EXISTS signal_cache (
                stock_code TEXT NOT NULL,
                date TEXT NOT NULL,
                PRIMARY KEY (stock_code, date)
            )
        ''')
        conn.commit()

    def _create_signal_table(self, signal_id: int):
        """创建信号缓存表"""
        conn = self._get_db_connection()
        table_name = f"signal_{signal_id}"
        try:
            # 创建表（如果不存在）
            conn.execute(f'''
                        CREATE TABLE IF NOT EXISTS {table_name} (
                            stock_code TEXT NOT NULL,
                            date TEXT NOT NULL,
                            signal_value REAL,
                            signal_strength REAL,
                            last_updated TEXT NOT NULL,
                            stock_name TEXT,
                            industry TEXT,
                            rating TEXT,
                            close_price REAL,
                            change_pct REAL,
                            PRIMARY KEY (stock_code, date)
                        )
                    ''')
            conn.commit()
        except Exception as e:
            logging.error(f"创建表 {table_name} 失败: {e}")

    def calculate_technical_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算必要的技术指标（精简版），成交量单位统一为万手
        :param df: 包含原始交易数据的DataFrame
        :return: 包含技术指标的DataFrame
        """
        if df.empty:
            return df

        # 确保数据按日期排序
        df = df.sort_values('date').reset_index(drop=True)

        # 创建副本避免修改原始数据
        df = df.copy()

        # 将成交量转换为万手单位（除以10000）并保留一位小数
        df['volume'] = (df['volume'] / 10000).round(1)

        # 1. 计算移动平均线
        # 计算移动平均线（保留一位小数）
        for window in [5, 10, 20, 60, 120]:
            # 计算收盘价移动平均线并四舍五入到1位小数
            df[f'ma{window}'] = df['close'].rolling(window=window, min_periods=1).mean().round(1)

            # 计算成交量移动平均线并四舍五入到1位小数
            df[f'vol_ma{window}'] = df['volume'].rolling(window=window, min_periods=1).mean().round(1)

        # 2. 计算RSI
        df['price_rsi'] = self._calculate_rsi(df['close'], period=14)
        df['vol_rsi'] = self._calculate_rsi(df['volume'], period=14)

        # 3. 计算MACD（成交量）
        exp12 = df['volume'].ewm(span=12, adjust=False, min_periods=1).mean()
        exp26 = df['volume'].ewm(span=26, adjust=False, min_periods=1).mean()
        df['vol_macd'] = exp12 - exp26
        df['vol_signal'] = df['vol_macd'].ewm(span=9, adjust=False, min_periods=1).mean()
        df['vol_hist'] = df['vol_macd'] - df['vol_signal']

        # 4. 计算布林带（成交量）
        df['vol_boll_middle'] = df['volume'].rolling(20, min_periods=1).mean()
        df['vol_boll_std'] = df['volume'].rolling(20, min_periods=1).std().fillna(0)
        df['vol_boll_upper'] = df['vol_boll_middle'] + 2 * df['vol_boll_std']
        df['vol_boll_lower'] = df['vol_boll_middle'] - 2 * df['vol_boll_std']

        # 5. 计算KDJ（成交量）
        low_min = df['volume'].rolling(window=9, min_periods=1).min()
        high_max = df['volume'].rolling(window=9, min_periods=1).max()
        rsv = ((df['volume'] - low_min) / (high_max - low_min + 1e-8)) * 100

        # 初始化KDJ值
        k = np.zeros(len(df))
        d = np.zeros(len(df))
        j = np.zeros(len(df))

        # 初始值
        k[0] = 50
        d[0] = 50
        j[0] = 50

        # 计算KDJ
        for i in range(1, len(df)):
            k[i] = (2 / 3) * k[i - 1] + (1 / 3) * rsv.iloc[i]
            d[i] = (2 / 3) * d[i - 1] + (1 / 3) * k[i]
            j[i] = 3 * k[i] - 2 * d[i]

        df['vol_k'] = k
        df['vol_d'] = d
        df['vol_j'] = j

        return df

    def _calculate_rsi(self, series: pd.Series, period: int = 14) -> pd.Series:
        """计算相对强弱指数(RSI)"""
        delta = series.diff()
        gain = (delta.where(delta > 0, 0)).fillna(0)
        loss = (-delta.where(delta < 0, 0)).fillna(0)

        # 使用指数移动平均
        avg_gain = gain.ewm(alpha=1 / period, min_periods=period).mean()
        avg_loss = loss.ewm(alpha=1 / period, min_periods=period).mean()

        # 计算RS
        rs = avg_gain / (avg_loss + 1e-8)  # 避免除零
        return 100 - (100 / (1 + rs))

    def _get_cached_signal(self, signal_id: int, stock_code: str, date: str) -> Optional[Dict]:
        """从缓存中获取信号数据"""
        conn = self._get_db_connection()
        table_name = f"signal_{signal_id}"
        try:
            cursor = conn.execute(
                f"SELECT signal_value, signal_strength FROM {table_name} WHERE stock_code = ? AND date = ?",
                (stock_code, date))
            row = cursor.fetchone()
            if row:
                return {
                    'signal_value': row[0],
                    'signal_strength': row[1]
                }
        except Exception as e:
            logging.error(f"读取缓存失败: {table_name} {stock_code}-{date}: {e}")
        return None

    def _save_signal_to_cache(self, signal_id: int, stock_code: str, date: str,
                              signal_value: float, signal_strength: float,
                              stock_name: str = "未知", industry: str = "未知",
                              rating: str = "N/A", close_price: float = 0.0,
                              change_pct: float = 0.0):
        """保存信号数据到缓存（更新版，包含新字段）"""
        conn = self._get_db_connection()
        table_name = f"signal_{signal_id}"
        current_time = datetime.now().strftime('%Y%m%d%H%M%S')
        try:
            conn.execute(
                f"""
                INSERT OR REPLACE INTO {table_name} 
                (stock_code, date, signal_value, signal_strength, last_updated,
                 stock_name, industry, rating, close_price, change_pct) 
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (stock_code, date, signal_value, signal_strength, current_time,
                 stock_name, industry, rating, close_price, change_pct))
            conn.commit()
        except Exception as e:
            logging.error(f"保存缓存失败: {table_name} {stock_code}-{date}: {e}")

    def generate_signals(self, df: Union[pd.DataFrame, List[pd.DataFrame]]) -> Dict[str, Any]:
        """
        生成交易信号，支持单股票或多股票并行处理
        :param df: 单个股票的DataFrame或股票数据列表
        :return: 包含所有信号的字典
        """
        if isinstance(df, list):
            # 多股票并行处理
            return self._generate_signals_parallel(df)
        else:
            # 单股票处理
            return self._generate_single_stock_signals(df)

    def _generate_signals_parallel(self, data_list: List[pd.DataFrame]) -> Dict[str, Any]:
        """并行处理多个股票数据"""
        results = {'signals_list': [], 'volume_signals': [], 'price_signals': [], 'other_signals': []}
        total_stocks = len(data_list)
        logging.info(f"开始并行处理 {total_stocks} 只股票的信号生成，使用 {self.max_workers} 个线程")

        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = {executor.submit(self._generate_single_stock_signals, df): df for df in data_list}

            for future in as_completed(futures):
                try:
                    result = future.result()
                    results['signals_list'].extend(result['signals_list'])
                    results['volume_signals'].extend(result['volume_signals'])
                    results['price_signals'].extend(result['price_signals'])
                    results['other_signals'].extend(result['other_signals'])
                except Exception as e:
                    df = futures[future]
                    stock_code = df['symbol'].iloc[0] if not df.empty else 'unknown'
                    logging.error(f"处理股票 {stock_code} 时出错: {str(e)}")

        logging.info(f"并行处理完成，共生成 {len(results['signals_list'])} 个信号")
        return results

    def _generate_single_stock_signals(self, df: pd.DataFrame) -> Dict[str, Any]:
        """处理单个股票的信号生成"""
        if df.empty:
            logging.warning("输入数据为空，无法生成信号")
            return {'signals_list': [], 'volume_signals': [], 'price_signals': [], 'other_signals': []}

        # 记录数据日期范围
        stock_code = df['symbol'].iloc[0] if 'symbol' in df.columns else 'unknown'

        # 确保日期是datetime类型
        if not pd.api.types.is_datetime64_any_dtype(df['date']):
            try:
                df['date'] = pd.to_datetime(df['date'])
            except Exception as e:
                logging.error(f"无法转换日期列: {e}")
                return {'signals_list': [], 'volume_signals': [], 'price_signals': [], 'other_signals': []}

        min_date = df['date'].min().strftime('%Y%m%d')
        max_date = df['date'].max().strftime('%Y%m%d')
        logging.info(f"生成股票 {stock_code} 信号: 数据范围 {min_date} 至 {max_date}, {len(df)} 条记录")

        # 计算技术指标（已包含成交量转换）
        signals = self.calculate_technical_indicators(df.copy())

        # 添加成交量列（已在calculate_technical_indicators中转换为万手）
        signals['volume_wan'] = signals['volume']  # 已经是万手单位

        # 生成信号
        signals = self._generate_basic_signals(signals)
        signals = self._generate_volume_signals(signals)

        # 创建信号输出并分类
        signal_output = self._create_signal_output(signals)

        logging.info(f"股票 {stock_code} 信号生成完成: 共发现 {len(signal_output['signals_list'])} 个信号")
        return signal_output

    def _generate_basic_signals(self, signals: pd.DataFrame) -> pd.DataFrame:
        """生成基础信号规则（成交量单位已统一为万手）"""
        # 信号1：成交量双均线条件
        signals = calculate_volume_signal(signals)

        # 信号2：量能趋势突破 - 改进版本
        vol_ma5_slope = signals['vol_ma5'].diff(2) > 0  # 2天内5日均量上涨
        signals['signal_vol_break'] = (
                (signals['vol_ma5'] > signals['vol_ma20']) &
                (signals['vol_ma5'].shift(1) <= signals['vol_ma20'].shift(1)) &
                vol_ma5_slope
        ).fillna(False)

        # 信号3：量价齐升
        signals['signal_volume_price'] = (
                (signals['volume'] > signals['vol_ma5'] * 1.1) &
                (signals['close'] > signals['open']) &
                (signals['close'] > signals['high'].shift(1))  # 修正：与前一日高点比较
        ).fillna(False)

        # 信号11：价格均线多头排列 - 改进版本
        signals['ma5_slope'] = signals['ma5'].diff(2) > 0  # 2天内5日均线上涨
        signals['signal_price'] = (
                (signals['ma5'] > signals['ma10']) &
                (signals['ma10'] > signals['ma20']) &
                (signals['close'] > signals['ma20']) &
                signals['ma5_slope']
        ).fillna(False)

        return signals

    def _generate_volume_signals(self, signals: pd.DataFrame) -> pd.DataFrame:
        """生成基于成交量指标的信号"""
        # 信号4：成交量MACD金叉
        signals['signal_vol_macd'] = (
                (signals['vol_macd'] > signals['vol_signal']) &
                (signals['vol_macd'].shift(1) <= signals['vol_signal'].shift(1))
        ).fillna(False)

        # 信号5：成交量RSI超卖
        signals['signal_vol_rsi'] = (signals['vol_rsi'] < 30).fillna(False)

        # 信号6：成交量布林带收窄
        if 'vol_boll_upper' in signals.columns and 'vol_boll_lower' in signals.columns:
            boll_width = signals['vol_boll_upper'] - signals['vol_boll_lower']
            avg_boll_width = boll_width.rolling(20, min_periods=1).mean()
            signals['signal_vol_boll'] = (
                    (boll_width < avg_boll_width * 0.8) &
                    (signals['volume'] > signals['vol_boll_lower']) &
                    (signals['volume'] < signals['vol_boll_upper'])
            ).fillna(False)
        else:
            signals['signal_vol_boll'] = False

        # 信号7：成交量KDJ金叉
        if 'vol_k' in signals.columns and 'vol_d' in signals.columns:
            signals['signal_vol_kdj'] = (
                    (signals['vol_k'] > signals['vol_d']) &
                    (signals['vol_k'].shift(1) <= signals['vol_d'].shift(1)) &
                    (signals['vol_j'] > 20)
            ).fillna(False)
        else:
            signals['signal_vol_kdj'] = False

        # 信号8：低波动率信号 - 改进版本
        # 增加连续低成交量确认
        low_volume_days = (signals['volume'] < signals['vol_ma20'] * 0.8).rolling(window=3).sum()

        # 增加价格波动率确认（如果有ATR指标）
        if 'atr_14' in signals.columns:
            low_price_volatility = signals['atr_14'] < signals['atr_14'].rolling(window=20).mean()
        else:
            # 如果没有ATR指标，使用价格变动幅度作为替代
            price_range = (signals['high'] - signals['low']) / signals['close'].shift(1)
            low_price_volatility = price_range < price_range.rolling(window=20).mean()

        signals['signal_low_volatility'] = (
                (low_volume_days >= 2) &
                low_price_volatility
        ).fillna(False)

        return signals

    def _create_signal_output(self, signals: pd.DataFrame) -> Dict[str, List]:
        """创建信号输出结构并存储到数据库，分类返回成交量信号和其他信号"""
        all_signals = []
        volume_signals = []
        price_signals = []
        other_signals = []

        # 确保有股票代码
        stock_code = signals['symbol'].iloc[0] if 'symbol' in signals.columns else 'unknown'

        # 从 DataFrame 中获取股票基本信息（所有行相同）
        stock_name = signals['stock_name'].iloc[0] if 'stock_name' in signals.columns else '未知'
        industry = signals['industry'].iloc[0] if 'industry' in signals.columns else '未知'
        rating = signals['rating'].iloc[0] if 'rating' in signals.columns else 'N/A'

        # 收集所有信号触发信息
        for signal_col, signal_id in self.SIGNAL_MAP.items():
            if signal_col not in signals.columns:
                continue

            # 获取触发信号的日期
            trigger_mask = signals[signal_col].fillna(False)
            trigger_rows = signals[trigger_mask]

            # 为每个触发行创建详细记录并存储到数据库
            for idx, row in trigger_rows.iterrows():
                # 从当前行获取涨跌幅值
                change_pct = row.get('change_pct', 0.0)

                # 获取该日期所有相关指标值
                date_str = row['date'].strftime('%Y%m%d') if not isinstance(row['date'], str) else row['date']

                # 获取最低价格
                close_price = row.get('low', 0)

                record = {
                    'date': date_str,
                    'stock_code': stock_code,
                    'stock_name': stock_name,  # 添加股票名称
                    'industry': industry,  # 添加行业
                    'rating': rating,  # 添加评级
                    'signal_id': signal_id,
                    'signal_name': signal_col,
                    'close_price': round(close_price, 2),
                    'volume': row.get('volume_wan', 0),  # 已经是万手单位
                    'price_rsi': round(row.get('price_rsi', 50), 1),
                    'vol_rsi': round(row.get('vol_rsi', 50), 1),
                    'vol_macd': round(row.get('vol_macd', 0), 3),
                    'vol_k': round(row.get('vol_k', 50), 1),
                    'vol_d': round(row.get('vol_d', 50), 1),
                    'change_pct': change_pct,  # 使用从行中获取的涨跌幅值
                    'description': self.SIGNAL_DESCRIPTIONS.get(signal_id, "未知信号")
                }
                all_signals.append(record)

                # 分类信号
                if signal_id <= 8:  # 成交量信号
                    volume_signals.append(record)
                elif 11 <= signal_id <= 14:  # 价格信号
                    price_signals.append(record)
                else:  # 其他信号
                    other_signals.append(record)

                # 存储到信号缓存数据库（添加新字段）
                self._save_signal_to_cache(
                    signal_id=signal_id,
                    stock_code=stock_code,
                    date=date_str,
                    signal_value=1,
                    signal_strength=1.0,
                    # 添加额外字段
                    stock_name=stock_name,
                    industry=industry,
                    rating=rating,
                    close_price=round(close_price, 2),
                    change_pct=change_pct
                )

        # 按信号ID排序
        all_signals.sort(key=lambda x: (x['date'], x['signal_id']))
        volume_signals.sort(key=lambda x: (x['date'], x['signal_id']))
        price_signals.sort(key=lambda x: (x['date'], x['signal_id']))
        other_signals.sort(key=lambda x: (x['date'], x['signal_id']))

        return {
            'signals_list': all_signals,
            'volume_signals': volume_signals,
            'price_signals': price_signals,
            'other_signals': other_signals
        }

    def close_connections(self):
        """关闭所有数据库连接"""
        # 关闭线程本地连接
        if hasattr(self, 'local'):
            try:
                if hasattr(self.local, 'conn') and self.local.conn:
                    self.local.conn.close()
            except:
                pass

        # 重置连接
        self.local = threading.local()

    def __del__(self):
        self.close_connections()

    def query_signal_data_all(self, signal_id: int, start_date: str = None, end_date: str = None,
                              stock_codes: List[str] = None) -> pd.DataFrame:
        """
        查询指定信号表中的所有数据，支持多种条件筛选
        不进行第二步筛选（每个股票代码只返回日期最接近结束时间的记录）

        :param signal_id: 信号ID
        :param start_date: 查询开始日期 (YYYYMMDD)，可选
        :param end_date: 查询结束日期 (YYYYMMDD)，可选
        :param stock_codes: 股票代码列表，可选
        :return: 包含查询结果的DataFrame
        """
        conn = self._get_db_connection()
        table_name = f"signal_{signal_id}"

        # 验证日期格式（如果提供了日期）
        if start_date and (len(start_date) != 8 or not start_date.isdigit()):
            logging.error("开始日期格式错误，请使用YYYYMMDD格式")
            return pd.DataFrame()

        if end_date and (len(end_date) != 8 or not end_date.isdigit()):
            logging.error("结束日期格式错误，请使用YYYYMMDD格式")
            return pd.DataFrame()

        try:
            # 构建查询语句
            query_parts = ["SELECT * FROM {table_name}"]
            params = []

            # 添加日期条件
            conditions = []
            if start_date:
                conditions.append("date >= ?")
                params.append(start_date)
            if end_date:
                conditions.append("date <= ?")
                params.append(end_date)

            # 添加股票代码条件
            if stock_codes:
                # 清理股票代码格式
                clean_codes = [re.sub(r'\D', '', code) for code in stock_codes]
                placeholders = ','.join(['?' for _ in clean_codes])
                conditions.append(f"stock_code IN ({placeholders})")
                params.extend(clean_codes)

            # 添加WHERE子句（如果有条件）
            if conditions:
                query_parts.append("WHERE " + " AND ".join(conditions))

            # 构建完整查询
            query = " ".join(query_parts).format(table_name=table_name)

            # 执行查询
            result = pd.read_sql_query(query, conn, params=params if params else None)

            if result.empty:
                date_range = f"{start_date} 到 {end_date}" if start_date or end_date else "所有日期"
                stock_info = f"，股票代码: {stock_codes}" if stock_codes else ""
                logging.info(f"信号表 {table_name} 在 {date_range}{stock_info} 范围内没有数据")

            return result

        except Exception as e:
            logging.error(f"查询信号表 {table_name} 失败: {e}")
            return pd.DataFrame()

    def query_signal_data_latest(self, signal_id: int, start_date: str = None, end_date: str = None,
                                 stock_codes: List[str] = None) -> pd.DataFrame:
        """
        查询指定信号表中的数据，支持多种条件筛选
        对于每个股票代码，只返回日期最接近结束时间的记录

        :param signal_id: 信号ID
        :param start_date: 查询开始日期 (YYYYMMDD)，可选
        :param end_date: 查询结束日期 (YYYYMMDD)，可选
        :param stock_codes: 股票代码列表，可选
        :return: 包含查询结果的DataFrame
        """
        # 第一步：获取所有数据
        all_data = self.query_signal_data_all(signal_id, start_date, end_date, stock_codes)

        if all_data.empty:
            return pd.DataFrame()

        # 第二步：对每个股票代码，筛选日期最接近结束时间的记录
        # 首先确保日期列是datetime类型
        all_data['date'] = pd.to_datetime(all_data['date'], format='%Y%m%d')

        # 确定结束日期（如果未提供，则使用数据中的最大日期）
        if end_date:
            end_date_dt = pd.to_datetime(end_date, format='%Y%m%d')
        else:
            end_date_dt = all_data['date'].max()

        # 计算每个记录与结束日期的绝对差值
        all_data['date_diff'] = (all_data['date'] - end_date_dt).abs()

        # 按股票代码分组，并取日期差值最小的记录
        idx = all_data.groupby('stock_code')['date_diff'].idxmin()
        result = all_data.loc[idx].drop(columns=['date_diff'])

        # 将日期列转换回字符串格式（YYYYMMDD）
        result['date'] = result['date'].dt.strftime('%Y%m%d')

        # 按日期和股票代码排序
        result = result.sort_values(by=['date', 'stock_code'])

        # 重置索引并删除旧索引列
        result = result.reset_index(drop=True)

        return result


# 测试代码
if __name__ == "__main__":
    logging.info("启动信号生成器测试...")

    # 创建信号生成器实例
    generator = SignalGenerator(max_workers=2)

    # 测试查询功能
    print("\n===== 测试信号数据查询功能 =====")
    # 查询信号1在指定时间范围内的数据

    all_data = generator.query_signal_data_all(
        signal_id=1,
        start_date="20250101",
        end_date="20251231"
    )
    print(f"获取到指定股票的所有数据 {len(all_data)} 条记录")

    if not all_data.empty:
        # 过滤掉以 '300'（创业板）、'688'（科创板）和 '920'（新三板） 和 '301' （北交所） 开头的股票代码
        filtered_data = all_data[~all_data['stock_code'].str.startswith(('30', '68', '92'))]

        if not filtered_data.empty:
            print(f"查询到 {len(filtered_data)} 条记录（去掉创业板和科创板股票）:")
            for index, row in filtered_data.iterrows():
                print(
                    f"股票代码: {row['stock_code']}, 日期: {row['date']}, 信号值: {row['signal_value']}, 最低价: {row['close_price']}, 涨跌幅: {row['change_pct']}")
        else:
            print("未查询到符合条件的数据（去掉创业板和科创板后）")
    else:
        print("未查询到符合条件的数据")


    # 创建测试数据
    # def create_test_data(stock_code: str, days: int = 10):
    #     """创建测试数据"""
    #     # 使用YYYYMMDD格式的日期字符串
    #     dates = pd.date_range(start='2025-06-11', periods=days).strftime('%Y%m%d').tolist()
    #
    #     data = {
    #         'date': pd.to_datetime(dates),  # 转换为datetime对象
    #         'symbol': [stock_code] * days,
    #         'open': [11.46, 11.56, 11.68, 11.57, 11.79, 11.78, 11.76, 11.70, 11.81, 11.90][:days],
    #         'close': [11.49, 11.68, 11.58, 11.79, 11.76, 11.77, 11.70, 11.84, 11.93, 11.93][:days],
    #         'high': [11.57, 11.70, 11.74, 11.79, 11.88, 11.81, 11.78, 11.86, 11.97, 12.01][:days],
    #         'low': [11.43, 11.52, 11.56, 11.53, 11.71, 11.66, 11.65, 11.67, 11.67, 11.84][:days],
    #         'volume': [1094667, 1290220, 1095494, 1174438, 890027, 683521, 832247, 1309198, 1389170, 1308525][:days],
    #     }
    #     return pd.DataFrame(data)
    #
    #
    # # 创建多个股票的测试数据
    # test_data = [
    #     create_test_data('000001'),
    #     create_test_data('000002'),
    #     create_test_data('600000'),
    #     create_test_data('600001'),
    # ]

    # # 测试单股票模式
    # print("\n===== 测试单股票模式 =====")
    # single_result = generator.generate_signals(test_data[0])
    # print(f"单股票模式生成 {len(single_result['signals_list'])} 个信号")
    # print(f"其中成交量信号: {len(single_result['volume_signals'])} 个")
    # print(f"价格信号: {len(single_result['price_signals'])} 个")
    # print(f"其他信号: {len(single_result['other_signals'])} 个")
    #
    # # 测试多股票并行模式
    # print("\n===== 测试多股票并行模式 =====")
    # multi_result = generator.generate_signals(test_data)
    # print(f"多股票并行模式生成 {len(multi_result['signals_list'])} 个信号")
    # print(f"其中成交量信号: {len(multi_result['volume_signals'])} 个")
    # print(f"价格信号: {len(multi_result['price_signals'])} 个")
    # print(f"其他信号: {len(multi_result['other_signals'])} 个")

    # # 打印成交量信号
    # print("\n成交量信号:")
    # for signal in multi_result['volume_signals']:
    #     print(f"[{signal['date']}] {signal['stock_code']} - {signal['signal_name']} ({signal['signal_id']})")
    #     print(f"  成交量: {signal['volume']} 万手")
    #     print(f"  描述: {signal['description']}")
    #
    # # 打印价格信号
    # print("\n价格信号:")
    # for signal in multi_result['price_signals']:
    #     print(f"[{signal['date']}] {signal['stock_code']} - {signal['signal_name']} ({signal['signal_id']})")
    #     print(f"  收盘价: {signal['close_price']}")
    #     print(f"  描述: {signal['description']}")
    #
    # # 打印其他信号
    # print("\n其他信号:")
    # for signal in multi_result['other_signals']:
    #     print(f"[{signal['date']}] {signal['stock_code']} - {signal['signal_name']} ({signal['signal_id']})")
    #     print(f"  描述: {signal['description']}")
    #
    # logging.info("信号生成器测试完成")
    # generator.close_connections()